import numpy as np
import torch
import torch.nn as nn
import torch.utils.data as Data
import random


# Hyper Parameters
DATASET_PATH = './datasets/ML-1m/ml-1m.txt'
USER_CNT = 6040
ITEM_CNT = 3416
MAX_POS = 200

RANDOM_MASK_CNT = 16
MASK_RATE = 0.2
SINGLE_MASK_RATE = 0.2


# DataSet
class BertDataset(Data.Dataset):
    def __init__(self, dataset_len):
        self.dataset_len = dataset_len

    def __getitem__(self, index):
        return generate_single_data()

    def __len__(self):
        return self.dataset_len


# Data Class
class Rate_Info:
    def __init__(self, user_id, item_id, rate_value, rate_time):
        self.user_id: int = user_id
        self.item_id: int = item_id
        self.rate_value: float = rate_value
        self.rate_time: float = rate_time


# Global Variables
rate_info = []
rate_matrix = torch.zeros((USER_CNT, ITEM_CNT + 2), dtype=torch.float)
user_matrix = None
user_rate_cnt = None

user_latent = None
item_latent = None

# Item Popular
item_popular = torch.zeros((ITEM_CNT + 2), dtype=torch.float)


# Load the data from file.
def load_from_file():
    print('Loading data...')
    fin = open(DATASET_PATH, 'rt')
    for line in fin.readlines():
        user_id, item_id = line.split(' ')
        user_id = int(user_id)-1
        item_id = int(item_id)
        rate_value = 1.0
        rate_time = 1.0
        rate_info.append(Rate_Info(user_id, item_id, rate_value, rate_time))
        rate_matrix[user_id, item_id] = rate_value
    print('Finished.')


# PMF Algorithm
def PMF(K: int, LR=0.0001, alpha=0.01, EPOCH=1000):
    global user_latent
    global item_latent

    try:
        user_latent = torch.load('./model_dat/user_latent.pkl')
        item_latent = torch.load('./model_dat/item_latent.pkl')

        print('PMF result loaded.')
    except Exception as e:
        print(e)
        print('Load PMF result failed, start training PMF.')

        user_latent = torch.randn(
            [USER_CNT, K], dtype=torch.float, requires_grad=True)
        item_latent = torch.randn(
            [ITEM_CNT, K], dtype=torch.float, requires_grad=True)
        I_matrix: torch.Tensor = (
            rate_matrix > torch.zeros_like(rate_matrix)).type(torch.float)

        for epoch in range(EPOCH):
            u_l2 = torch.sum(user_latent ** 2)
            i_l2 = torch.sum(item_latent ** 2)

            loss = torch.sum(I_matrix.mul(
                (rate_matrix - user_latent.mm(item_latent.t())) ** 2)) + alpha * (u_l2 + i_l2)

            user_latent.grad = torch.zeros_like(user_latent)
            item_latent.grad = torch.zeros_like(item_latent)
            loss.backward()

            user_latent.data -= LR * user_latent.grad.data
            item_latent.data -= LR * item_latent.grad.data

            print(f'EPOCH:{epoch},LOSS:{loss.item()}')

        torch.save(user_latent, './model_dat/user_latent.pkl')
        torch.save(item_latent, './model_dat/item_latent.pkl')
        print('PMF result saved.')

    finally:
        # Set requires_grad to False.
        item_latent.requires_grad = False
        user_latent.requires_grad = False

        print('PMF finished.')


# Load Data
def load_data():
    print('Start loading the data...')

    # Initialize User_Matrix.
    global user_rate_cnt, user_matrix, rate_info, item_popular

    try:
        user_rate_cnt, user_matrix, item_popular = torch.load(
            './model_dat/user_matrix_bert.pkl')
    except Exception as e:
        print(e)
        print('Load from file failed.')

        # Load the data from the datasets files.
        load_from_file()

        user_matrix = torch.zeros((USER_CNT, MAX_POS), dtype=torch.long)
        user_rate_cnt = torch.zeros(USER_CNT, dtype=torch.long)

        # Sort the rate items by time.
        # rate_info = sorted(rate_info, key=lambda x: x.rate_time, reverse=True)
        rate_info.reverse()

        print('Loading into tensor...')

        rate_info_len = len(rate_info)
        for i, rate in enumerate(rate_info):
            # Write into the item popular.
            item_popular[rate.item_id] += 1.0

            if user_rate_cnt[rate.user_id] < MAX_POS:
                user_rate_cnt[rate.user_id] += 1
                user_matrix[rate.user_id, -user_rate_cnt[rate.user_id]
                            ] = rate.item_id

            if i % 20000 == 0:
                print(f'{i}/{rate_info_len} finished.')

        # Softmax the item popular.
        item_popular[0] = 0.0

        # Save the pkl.
        torch.save((user_rate_cnt, user_matrix, item_popular),
                   './model_dat/user_matrix_bert.pkl')

    print('Data loaded successfully.')


# Generate A Single Data
def generate_single_data():
    # Random user id.
    cur_user_id = random.randint(0, USER_CNT - 1)
    cur_cnt = user_rate_cnt[cur_user_id].item()

    cur_src = torch.zeros([MAX_POS], dtype=torch.long)
    cur_src[:cur_cnt - 1] = user_matrix[cur_user_id, -cur_cnt:-1]

    cur_data_x = torch.zeros(
        [MAX_POS], dtype=torch.long)
    cur_padding = torch.zeros([MAX_POS], dtype=torch.bool)
    cur_data_y = torch.zeros([MAX_POS], dtype=torch.long)
    cur_masked = torch.zeros([MAX_POS], dtype=torch.long)

    # Random Mask
    if random.random() > SINGLE_MASK_RATE:
        cur_padding = torch.cat((torch.zeros(
            cur_cnt - 1), torch.ones(MAX_POS - cur_cnt + 1)), dim=0).bool()

        cur_masked = (torch.rand(MAX_POS) < MASK_RATE).type(torch.long)
        # cur_masked[cur_cnt - 2] = 1
        cur_data_y = cur_src
        cur_data_x = cur_src.masked_fill(cur_masked, ITEM_CNT + 1)
    else:
        # Single Mask
        seq_len = cur_cnt - 1

        cur_padding = torch.cat((torch.zeros(
            seq_len), torch.ones(MAX_POS - seq_len)), dim=0).bool()

        cur_masked[seq_len - 1] = 1
        cur_data_y = cur_src
        cur_data_x = cur_src.masked_fill(cur_masked, ITEM_CNT+1)

    return cur_data_x, cur_data_y, cur_padding, cur_user_id, cur_cnt - 1, cur_masked


# Generate Data
def generate_data_user(cur_user_id):
    cur_cnt = user_rate_cnt[cur_user_id].item()

    # Check Special User
    if cur_cnt < 3:
        return None, None, None

    cur_src = torch.zeros([MAX_POS], dtype=torch.long)
    cur_src[:cur_cnt - 1] = user_matrix[cur_user_id, -cur_cnt:-1]

    cur_data_x = torch.zeros(
        [RANDOM_MASK_CNT, MAX_POS], dtype=torch.long)
    cur_padding = torch.zeros([RANDOM_MASK_CNT, MAX_POS], dtype=torch.bool)
    cur_data_y = torch.zeros([RANDOM_MASK_CNT, MAX_POS], dtype=torch.long)
    cur_masked = torch.zeros([RANDOM_MASK_CNT, MAX_POS], dtype=torch.long)

    for i in range(RANDOM_MASK_CNT):
        # Random Mask
        if i < RANDOM_MASK_CNT - 1:
            cur_padding[i] = torch.cat((torch.zeros(
                cur_cnt - 1), torch.ones(MAX_POS - cur_cnt + 1)), dim=0).bool()

            cur_masked[i] = (torch.rand(MAX_POS) < MASK_RATE).type(torch.long)
            # cur_masked[i, cur_cnt - 2] = 1
            cur_data_y[i] = cur_src
            cur_data_x[i] = cur_src.masked_fill(cur_masked[i], ITEM_CNT+1)
        else:
            # Single Mask
            # seq_len = random.randint(cur_cnt // 2, cur_cnt - 1)
            seq_len = cur_cnt - 1

            cur_padding[i] = torch.cat((torch.zeros(
                seq_len), torch.ones(MAX_POS - seq_len)), dim=0).bool()

            cur_masked[i, seq_len - 1] = 1
            cur_data_y[i] = cur_src
            cur_data_x[i] = cur_src.masked_fill(cur_masked[i], ITEM_CNT+1)

    return cur_data_x, cur_data_y, cur_padding, cur_masked


def generate_data(step):
    '''
    data_x: Input tensor, masked.
    data_y: Output tensor, not masked.
    train_padding: Data padding.
    train_masked: Data mask.
    '''

    print('Generating user data...')
    cur_user_id = 0
    while cur_user_id < USER_CNT:
        train_data_x = torch.empty([0, MAX_POS], dtype=torch.long)
        train_data_y = torch.empty([0, MAX_POS], dtype=torch.long)
        train_padding = torch.empty([0, MAX_POS], dtype=torch.bool)
        train_masked = torch.empty([0, MAX_POS], dtype=torch.long)

        end_user_id = min(cur_user_id + step, USER_CNT)
        while cur_user_id < end_user_id:
            cur_data_x, cur_data_y, cur_padding, cur_masked = generate_data_user(
                cur_user_id)
            if cur_data_x is not None and cur_data_y is not None:
                train_data_x = torch.cat((train_data_x, cur_data_x), dim=0)
                train_data_y = torch.cat((train_data_y, cur_data_y), dim=0)
                train_padding = torch.cat((train_padding, cur_padding), dim=0)
                train_masked = torch.cat((train_masked, cur_masked), dim=0)

            cur_user_id += 1

        # Yield a value.
        dataset = Data.TensorDataset(
            train_data_x, train_data_y, train_padding, train_masked)
        print('Finished.')
        yield dataset
    return


def generate_test_data():
    print('Generating test data...')
    test_data_x = torch.zeros([USER_CNT, MAX_POS], dtype=torch.long)
    test_data_y = torch.zeros([USER_CNT, MAX_POS], dtype=torch.long)
    test_padding = torch.zeros([USER_CNT, MAX_POS], dtype=torch.bool)
    test_masked = torch.zeros([USER_CNT, MAX_POS], dtype=torch.long)
    test_user_id = torch.zeros([USER_CNT], dtype=torch.long)
    test_cur_cnt = torch.zeros([USER_CNT], dtype=torch.long)

    for user_id in range(USER_CNT):
        cur_cnt = user_rate_cnt[user_id]

        test_data_x[user_id, :cur_cnt - 1] = user_matrix[user_id, -cur_cnt:-1]
        test_data_x[user_id, cur_cnt - 1] = ITEM_CNT + 1
        test_masked[user_id, cur_cnt - 1] = 1

        test_padding[user_id] = torch.cat((torch.zeros(
            cur_cnt), torch.ones(MAX_POS - cur_cnt)), dim=0).bool()
        test_data_y[user_id, cur_cnt - 1] = user_matrix[user_id, -1]
        test_user_id[user_id] = user_id
        test_cur_cnt[user_id] = cur_cnt

    dataset = Data.TensorDataset(
        test_data_x, test_data_y, test_padding, test_user_id, test_cur_cnt, test_masked)
    print('Test data generate successfully.')
    return dataset


def generate_train_dataset():
    data_x = torch.zeros(
        [RANDOM_MASK_CNT*USER_CNT, MAX_POS], dtype=torch.long)
    padding = torch.zeros(
        [RANDOM_MASK_CNT*USER_CNT, MAX_POS], dtype=torch.bool)
    data_y = torch.zeros([RANDOM_MASK_CNT*USER_CNT, MAX_POS], dtype=torch.long)
    masked = torch.zeros([RANDOM_MASK_CNT*USER_CNT, MAX_POS], dtype=torch.long)

    for user_id in range(USER_CNT):
        cur_data_x, cur_data_y, cur_padding, cur_masked = generate_data_user(
            user_id)
        data_x[user_id*RANDOM_MASK_CNT:(user_id+1) *
               RANDOM_MASK_CNT] = cur_data_x
        data_y[user_id*RANDOM_MASK_CNT:(user_id+1) *
               RANDOM_MASK_CNT] = cur_data_y
        padding[user_id*RANDOM_MASK_CNT:(user_id+1) *
                RANDOM_MASK_CNT] = cur_padding
        masked[user_id*RANDOM_MASK_CNT:(user_id+1) *
               RANDOM_MASK_CNT] = cur_masked

        if user_id % 1000 == 0:
            print(f'{user_id}/{USER_CNT} finished.')

    return Data.TensorDataset(data_x, data_y, padding, masked)


if __name__ == '__main__':
    load_from_file()
    PMF(K=64)
